/******************************************************************************
	Cobalt - utility libraries for building XNA games
	Copyright (C) 2008 Philip Djonov

	This program is free software; you can redistribute it and/or modify it
	under the terms of the GNU General Public License as published by the Free
	Software Foundation; either version 2 of the License, or (at your option)
	any later version.

	This program is distributed in the hope that it will be useful, but
	WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
	or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
	for more details.

	You should have received a copy of the GNU General Public License along
	with this program; if not, write to the
	
		Free Software Foundation, Inc.
		51 Franklin Street, Fifth Floor
		Boston, MA  02110-1301, USA.
******************************************************************************/

/*****
	Effect configuration parameters
*****/

float BloomThresh;

float3 BloomTint;
float BloomSaturation;

float3 BaseTint;
float BaseSaturation;

float3 DofRange;
float2 DofDefocus;

/*****
	Post-process driver interface.
*****/

float2 FullPixelSize, BlurPixelSize;
sampler FullImage : register( s0 );
sampler BlurImage : register( s1 );
sampler DepthImage : register( s2 );

#define GAUSS_TAP_COUNT 15
float4 GaussianTaps[(GAUSS_TAP_COUNT - 1) / 2];

#define POISSON_TAP_COUNT 8
float4 PoissonTaps[POISSON_TAP_COUNT - 1];

/*****
	Constantastic
*****/

#define DOF_COC_SIZE 5
#define DOF_RAD_SCALE 0.4

/*****
	Helper functions.
*****/

float luminance( float3 cl )
{
	return dot( cl, float3( 0.3, 0.59, 0.11 ) );
}

float3 adjust_saturation( float3 cl, float sat )
{
	float lum = luminance( cl );
	return lerp( lum, cl, sat );
}

/*****
	Vertex shader (common to all)
*****/

float4 post_proc_vs(
	float2 iPos : POSITION,
	float2 iTc : TEXCOORD0,

	out float2 oTc : TEXCOORD0
	) : POSITION
{
	oTc = iTc;
	return float4( iPos, 0, 1 );
}

/*****
	Downscale & Blur
*****/

float4 downscale_ps( float2 iTc : TEXCOORD0 ) : COLOR
{
	float3 cl;

	cl = tex2D( FullImage, iTc + float2( FullPixelSize.x, FullPixelSize.y ) ).rgb;
	cl += tex2D( FullImage, iTc + float2( FullPixelSize.x, -FullPixelSize.y ) ).rgb;
	cl += tex2D( FullImage, iTc + float2( -FullPixelSize.x, -FullPixelSize.y ) ).rgb;
	cl += tex2D( FullImage, iTc + float2( -FullPixelSize.x, FullPixelSize.y ) ).rgb;

	cl *= 0.25;

	return float4( cl, 0 );
}

float4 gauss_blur_ps( float2 iTc : TEXCOORD0 ) : COLOR
{
	float3 cl = tex2D( BlurImage, iTc ).rgb * GaussianTaps[0].w; //center tap
	for( int i = 0; i < (GAUSS_TAP_COUNT - 1) / 2; i++ )
	{
		cl += tex2D( BlurImage, iTc + GaussianTaps[i].xy ).rgb * GaussianTaps[i].z;
		cl += tex2D( BlurImage, iTc - GaussianTaps[i].xy ).rgb * GaussianTaps[i].z;
	}
	return float4( cl, 0 );
}

/*****
	Dof & Bloom
*****/

float compute_dof_defocus( float depth )
{
	float dofDepth = clamp( depth - DofRange.y, DofRange.x - DofRange.y, DofRange.z - DofRange.y );
	return -min( DofDefocus.x * dofDepth / (DofRange.y - DofRange.x), 0 ) +
		max( DofDefocus.y * dofDepth / (DofRange.z - DofRange.y), 0 );
}

float3 sample_base_dof( float2 tc, out float defocus )
{
	float centerDepth = tex2D( DepthImage, tc ).r;

	defocus = compute_dof_defocus( centerDepth );

	float cocRadiusFull = defocus * DOF_COC_SIZE;
	float cocRadiusBlur = cocRadiusFull * DOF_RAD_SCALE;

	float3 centerColor = tex2D( FullImage, tc ).rgb;

	float3 blur = 0;
	float blurW = 1;

	for( int i = 0; i < POISSON_TAP_COUNT - 1; i++ )
	{
		float2 tcFull = tc + PoissonTaps[i].xy * cocRadiusFull;	
		float2 tcBlur = tc + PoissonTaps[i].zw * cocRadiusBlur;

		float3 tapFull = tex2D( FullImage, tcFull ).rgb;	
		float3 tapBlur = tex2D( BlurImage, tcBlur ).rgb;

		float tapDepth = tex2D( DepthImage, tcFull ).r;
		float tapDefocus = compute_dof_defocus( tapDepth );
		
		float3 tapCl = lerp( tapFull, tapBlur, tapDefocus );
		
		float contrib = centerDepth < tapDepth ?
			abs( tapDefocus * 2 - 1 ) : 1;
		
		blur += tapCl * contrib;
		blurW += contrib;
	}
	
	return lerp( centerColor, blur / blurW, defocus );
}

float4 combine_dof_and_bloom_ps( float2 iTc : TEXCOORD0 ) : COLOR
{
	float defocus;
	float3 base = sample_base_dof( iTc, defocus );
	float3 bloom = tex2D( BlurImage, iTc ).rgb * saturate( (1 - defocus) * 2 );

	//scale away dark areas in the bloom
	float lum = luminance( bloom );
	float adjLum = saturate( (lum - BloomThresh) / (1 - BloomThresh) );
	bloom *= adjLum / lum;

	//adjust saturation and tint/intensity
	bloom = adjust_saturation( bloom, BloomSaturation ) * BloomTint;
	base = adjust_saturation( base, BaseSaturation ) * BaseTint * (1 - saturate( bloom ));

	//add and return values
	return float4( base + bloom, 1 );
}

/*****
	Bloom
*****/

float4 combine_bloom_ps( float2 iTc : TEXCOORD0 ) : COLOR
{
	float3 base = tex2D( FullImage, iTc ).rgb;
	float3 bloom = tex2D( BlurImage, iTc ).rgb;

	//scale away dark areas in the bloom
	float lum = luminance( bloom );
	float adjLum = saturate( (lum - BloomThresh) / (1 - BloomThresh) );
	bloom *= adjLum / lum;

	//adjust saturation and tint/intensity
	bloom = adjust_saturation( bloom, BloomSaturation ) * BloomTint;
	base = adjust_saturation( base, BaseSaturation ) * BaseTint * (1 - saturate( bloom ));

	//add and return values
	return float4( base + bloom, 1 );
}

/*****
	Passthrough
*****/

float4 combine_passthrough( float2 iTc : TEXCOORD0 ) : COLOR
{
	return float4( tex2D( FullImage, iTc ).rgb, 1 );
}

/*****
	Techniques
*****/

technique DofAndBloom
{
	pass Downscale
	{
		VertexShader = compile vs_2_0 post_proc_vs();
		PixelShader = compile ps_2_0 downscale_ps();

		ZEnable = False;
		AlphaTestEnable = False;
		AlphaBlendEnable = False;
	}

	pass Blur
	{
		VertexShader = compile vs_2_0 post_proc_vs();
		PixelShader = compile ps_2_0 gauss_blur_ps();
	}

	pass Combine
	{
		VertexShader = compile vs_2_0 post_proc_vs();
		PixelShader = compile ps_2_b combine_dof_and_bloom_ps();
	}
}

technique Bloom
{
	pass Downscale
	{
		VertexShader = compile vs_2_0 post_proc_vs();
		PixelShader = compile ps_2_0 downscale_ps();

		ZEnable = False;
		AlphaTestEnable = False;
		AlphaBlendEnable = False;
	}

	pass Blur
	{
		VertexShader = compile vs_2_0 post_proc_vs();
		PixelShader = compile ps_2_0 gauss_blur_ps();
	}

	pass Combine
	{
		VertexShader = compile vs_2_0 post_proc_vs();
		PixelShader = compile ps_2_0 combine_bloom_ps();
	}
}

technique Passthrough
{
	pass Combine
	{
		ZEnable = False;
		AlphaTestEnable = False;
		AlphaBlendEnable = False;
	
		VertexShader = compile vs_2_0 post_proc_vs();		
		PixelShader = compile ps_2_0 combine_passthrough();
	}
}